#! /usr/bin/python
# -*- encoding: utf-8 -*-
import numpy
import torch
import torchaudio
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import Parameter
from models.ResNetBlocks import *
from torch import nn
from torch.nn import init

from matplotlib import pyplot as plt

class SpatialAttention(nn.Module):
    # 直接使用 7*7 的卷积核训练出一个权重图
    def __init__(self, kernel_size=7):
        super().__init__()
        self.conv = nn.Conv2d(1, 1, kernel_size=kernel_size, padding=kernel_size // 2)
        # self.conv = nn.Conv2d(1, 1, kernel_size=kernel_size, padding=(19, 100))
        # self.conv = nn.Conv2d(2, 1, kernel_size=kernel_size, padding=0)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # print("x1:", numpy.shape(x))
        # max_result, _ = torch.max(x, dim=1, keepdim=True)
        # avg_result = torch.mean(x, dim=1, keepdim=True)
        # result = torch.cat([max_result, avg_result], 1)
        # print("result", numpy.shape(result))
        output = self.conv(x)
        # print("requires_grad",self.conv(x).requires_grad)
        # print("output", numpy.shape(output))
        # print("output", output[0][0])
        output = self.sigmoid(output)
        # print("output", numpy.shape(output))
        # exit()
        return output


class CBAMBlock(nn.Module):

    def __init__(self, channel=512, reduction=16, kernel_size=7):
        super().__init__()
        # self.ca = ChannelAttention(channel=channel, reduction=reduction)
        self.sa = SpatialAttention(kernel_size=kernel_size)

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, x):
        b, c, _, _ = x.size()
        residual = x
        # print("sa", self.sa(x).requires_grad)
        out = x * self.sa(x)

        return out + residual

class FrequencyAttention(nn.Module):
    def __init__(self, Filter_Num=512, reduction=4):
        super().__init__()

        self.se1 = nn.Sequential(
            nn.Conv2d(Filter_Num, Filter_Num // reduction, 1, bias=False),
            nn.ReLU(),
            nn.Conv2d(Filter_Num // reduction, Filter_Num, 1, bias=False)
        )
        self.se2 = nn.Sequential(
            nn.Conv2d(Filter_Num, Filter_Num // reduction, 1, bias=False),
            nn.ReLU(),
            nn.Conv2d(Filter_Num // reduction, Filter_Num, 1, bias=False)
        )
        self.se3 = nn.Sequential(
            nn.Conv2d(Filter_Num, Filter_Num // reduction, 1, bias=False),
            nn.ReLU(),
            nn.Conv2d(Filter_Num // reduction, Filter_Num, 1, bias=False)
        )
        self.se4 = nn.Sequential(
            nn.Conv2d(Filter_Num, Filter_Num // reduction, 1, bias=False),
            nn.ReLU(),
            nn.Conv2d(Filter_Num // reduction, Filter_Num, 1, bias=False)
        )

        # self.conv = nn.Conv2d(2, 1, kernel_size=kernel_size, padding=kernel_size // 2)

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # print("x", numpy.shape(x))
        max_result_c, _ = torch.max(x, dim=1, keepdim=True)
        # print("output", numpy.shape(max_result_c))
        avg_result_c = torch.mean(x, dim=1, keepdim=True)
        # print("output", numpy.shape(avg_result_c))

        max_result_1, _ = torch.max(max_result_c, dim=3, keepdim=True)
        # print("max_result", numpy.shape(max_result_1))
        max_result_1 = self.se1(max_result_1.permute(0, 2, 1, 3))
        # print("max_result", numpy.shape(max_result_1))

        avg_result_1 = torch.mean(max_result_c, dim=3, keepdim=True)
        avg_result_1 = self.se2(avg_result_1.permute(0, 2, 1, 3))
        # print("avg_result", numpy.shape(avg_result_1))

        max_result_2, _ = torch.max(avg_result_c, dim=3, keepdim=True)
        max_result_2 = self.se3(max_result_2.permute(0, 2, 1, 3))
        # print("max_result", numpy.shape(max_result_2))

        avg_result_2 = torch.mean(avg_result_c, dim=3, keepdim=True)
        avg_result_2 = self.se4(avg_result_2.permute(0, 2, 1, 3))
        # print("avg_result", numpy.shape(avg_result_2))

        output = self.sigmoid(max_result_1 + avg_result_1 + max_result_2 + avg_result_2)
        # print("output", numpy.shape(output))
        output = output.permute(0, 2, 1, 3)

        # print("output", numpy.shape(output))
        # print("FA:", output[0, 0])
        # exit()
        return output


class FrequencyAttentionBlock(nn.Module):

    def __init__(self, Filter_Num=512, reduction=2):
        super().__init__()

        self.fa = FrequencyAttention(Filter_Num=Filter_Num, reduction=reduction)

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, x):
        b, c, _, _ = x.size()
        residual = x
        # print("sa", self.fa(x).requires_grad)

        out = x * self.fa(x)

        return out + residual

class ResNetSE(nn.Module):
    def __init__(self, block, layers, num_filters, nOut, encoder_type='SAP', n_mels=40, log_input=True, **kwargs):
        super(ResNetSE, self).__init__()

        print('Embedding size is %d, encoder %s.' % (nOut, encoder_type))

        # 定义一个CBAM
        self.cbam = CBAMBlock(channel=1, reduction=1, kernel_size=7)
        # 定义一个FA模块
        # self.fa = FrequencyAttentionBlock(Filter_Num=40)
        self.fa = FrequencyAttentionBlock(Filter_Num=80)

        self.inplanes = num_filters[0]
        self.encoder_type = encoder_type
        self.n_mels = n_mels
        self.log_input = log_input

        self.conv1 = nn.Conv2d(1, num_filters[0], kernel_size=7, stride=(2, 1), padding=3,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(num_filters[0])
        self.relu = nn.ReLU(inplace=True)

        self.layer1 = self._make_layer(block, num_filters[0], layers[0])
        self.layer2 = self._make_layer(block, num_filters[1], layers[1], stride=(2, 2))
        self.layer3 = self._make_layer(block, num_filters[2], layers[2], stride=(2, 2))
        self.layer4 = self._make_layer(block, num_filters[3], layers[3], stride=(1, 1))

        self.instancenorm = nn.InstanceNorm1d(n_mels)
        self.torchfb = torchaudio.transforms.MelSpectrogram(sample_rate=16000, n_fft=512, win_length=400,
                                                            hop_length=160, window_fn=torch.hamming_window,
                                                            n_mels=n_mels)

        if self.encoder_type == "SAP":
            self.sap_linear = nn.Linear(num_filters[3] * block.expansion, num_filters[3] * block.expansion)
            self.attention = self.new_parameter(num_filters[3] * block.expansion, 1)
            out_dim = num_filters[3] * block.expansion
        elif self.encoder_type == "ASP":
            self.sap_linear = nn.Linear(num_filters[3] * block.expansion, num_filters[3] * block.expansion)
            self.attention = self.new_parameter(num_filters[3] * block.expansion, 1)
            out_dim = num_filters[3] * block.expansion * 2
        else:
            raise ValueError('Undefined encoder')

        self.fc = nn.Linear(out_dim, nOut)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def new_parameter(self, *size):
        out = nn.Parameter(torch.FloatTensor(*size))
        nn.init.xavier_normal_(out)
        return out

    def forward(self, x):

        with torch.no_grad():
            with torch.cuda.amp.autocast(enabled=False):
                x = self.torchfb(x) + 1e-6
                if self.log_input: x = x.log()
                x = self.instancenorm(x).unsqueeze(1).detach()

        # plt.figure('谱图')
        # x = x.cpu()
        # plt.imshow(x[0][0].detach().numpy())
        # # plt.show()
        # plt.savefig('png/谱图', dpi=300)
        # x = x.cuda()
        # exit()

        x = self.fa(x)
        # print("fa x:", numpy.shape(x))
        # plt.figure('aa')
        # x = x.cpu()
        # plt.imshow(x[0][0].detach().numpy())
        # # plt.show()
        # plt.savefig('png/jiaquan谱图', dpi=300)
        # exit()


        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = torch.mean(x, dim=2, keepdim=True)
        # print("out_dim:", numpy.shape(x))
        if self.encoder_type == "SAP":
            x = x.permute(0, 3, 1, 2).squeeze(-1)
            # print("out_dim:", numpy.shape(x))
            h = torch.tanh(self.sap_linear(x))
            # print("out_dim:", numpy.shape(h))
            # exit()
            w = torch.matmul(h, self.attention).squeeze(dim=2)
            w = F.softmax(w, dim=1).view(x.size(0), x.size(1), 1)
            x = torch.sum(x * w, dim=1)

        elif self.encoder_type == "ASP":
            x = x.permute(0, 3, 1, 2).squeeze(-1)
            h = torch.tanh(self.sap_linear(x))
            # 通道上乘上权重
            w = torch.matmul(h, self.attention).squeeze(dim=2)
            w = F.softmax(w, dim=1).view(x.size(0), x.size(1), 1)
            mu = torch.sum(x * w, dim=1)
            rh = torch.sqrt((torch.sum((x ** 2) * w, dim=1) - mu ** 2).clamp(min=1e-5))
            x = torch.cat((mu, rh), 1)

        x = x.view(x.size()[0], -1)
        # print("out_dim:",numpy.shape(x))
        # exit()
        x = self.fc(x)
        # print("out_dim:", numpy.shape(x))
        # exit()
        return x


def MainModel(nOut=256, **kwargs):
    # Number of filters
    num_filters = [16, 32, 64, 128]
    model = ResNetSE(SEBasicBlock, [3, 4, 6, 3], num_filters, nOut, **kwargs)
    return model


